require(cem)
require(mgcv)
#install.packages("xtable")
require(xtable)
#install.packages("Matching")
require(Matching)
require(MatchIt)
#install.packages("stargazer")
require(stargazer)
require(mvtnorm)



##############################################
## Proportional plot
##############################################


prop.plot <- function(props,labels,cutpoint,x0,x1,lab.print=F,col.offset=1,other=T,col.type="gray"){
  props. <- props
  props <- props[labels!="OTHER"]
  labels <- labels[labels!="OTHER"]
  labels <- as.character(labels)[rev(order(props))]
  props <- props[rev(order(props))]
  n <- sum(props>cutpoint)
  if(col.type=="gray"){
  cols <- gray(rev(1:(n+col.offset)/(n+col.offset)))}
  if(col.type=="heat"){
    cols <- rev(heat.colors(n+col.offset))}
  if(col.type=="terrain"){
    cols <- rev(terrain.colors(n+col.offset))}
  if(col.type=="cm"){
    cols <- cm.colors(n+col.offset)
    cols <- rev(cols[1:n])
  }
  polygon(c(x0,x0,x1,x1),c(0,props[1],props[1],0),border=NA,col=gray(1:(n+1)/(n+1))[1])
  if(lab.print==T){text(x=x0,y=1,label=labels[1],pos=4,col=gray(rev(1:(n+1)/(n+1)))[1],cex=.5)}
  if(other==T){text(x=x0,y=sum(props),label="OTHER",pos=4,col=gray(rev(1:(n+1)/(n+1)))[n],cex=.5)}
  for(i in 2:n){ polygon(c(x0,x0,x1,x1),c(sum(props[1:(i-1)]),sum(props[1:(i-1)])+props[i],sum(props[1:(i-1)])+props[i],sum(props[1:(i-1)])),border=NA,col=gray(1:(n+1)/(n+1))[i])
  }
  if(lab.print==T){for(i in 2:n){text(x=x0,y=sum(props[1:(i-1)])+1,label=labels[i],pos=4,col=cols[i],cex=.5)}
  }
  polygon(c(x0,x1,x1,x0),c(0,0,sum(props.),sum(props.)),border="lightgray")
}




##############################################
## Proportional plot 2
##############################################


prop.plot2 <- function(props,labels,n=NULL,cutpoint=0,x0,x1,lab.print=F,cols=NULL,other=F){
  props. <- props
  props <- props[labels!="OTHER"]
  labels <- labels[labels!="OTHER"]
  labels <- as.character(labels)[rev(order(props))]
  props <- props[rev(order(props))]
  if(length(n)==0){n <- sum(props>cutpoint)}
  polygon(c(x0,x0,x1,x1),c(0,props[1],props[1],0),border=NA,col=gray(1:(n+1)/(n+1))[1])
  if(lab.print==T){text(x=x0,y=1,label=labels[1],pos=4,col=cols[1],cex=.5)}
  if(other==T){text(x=x0,y=sum(props.)-2,label="OTHER",pos=4,col="black",cex=.5)}
  for(i in 2:n){ polygon(c(x0,x0,x1,x1),c(sum(props[1:(i-1)]),sum(props[1:(i-1)])+props[i],sum(props[1:(i-1)])+props[i],sum(props[1:(i-1)])),border=NA,col=gray(1:(n+1)/(n+1))[i])
  }
  if(lab.print==T){for(i in 2:n){text(x=x0,y=sum(props[1:(i-1)])+1,label=labels[i],pos=4,col=cols[i],cex=.5)}
  }
  polygon(c(x0,x1,x1,x0),c(0,0,sum(props.),sum(props.)),border="lightgray")
}


##############################################
## Balance summary
##############################################


sum.balance <- function(X,labels=rightside){
PreT <- apply(X[X$Treat==1,rightside],2,mean)
PreC <- apply(X[X$Treat==0,rightside],2,mean)
PreSDM <- abs((apply(X[X$Treat==1,rightside],2,mean)-apply(X[X$Treat==0,rightside],2,mean))/apply(X[X$Treat==1,rightside],2,sd))
ks <- c()
tt <- c()
for(i in 1:length(rightside)){ks[i] <- ks.test(X[X$Treat==1,rightside[i]],X[X$Treat==0,rightside[i]])$p.value; tt[i] <- t.test(X[X$Treat==1,rightside[i]],X[X$Treat==0,rightside[i]])$p.value}
tab <- data.frame(Treated=PreT,Control=PreC,SDM=PreSDM,TTest=tt,KSTest=ks)
row.names(tab) <- labels
out <- list(Table=tab,Treated=sum(X$Treat==1),Control=sum(X$Treat==0),SDM=mean(PreSDM))
return(out)}



##############################################
## Risk ratio plot
##############################################

rr.plot <- function(pred.a,pred.b,pred.c,xlim,cex.axis=.9) {
par(mar=c(2,0,0,0))
plot(x=0,y=0,col=NA,xlim=xlim,ylim=c(-1,2),xaxt="n",bty="n")
rect(xleft=pred.a$rr.lo,xright=pred.a$rr.up,ybottom=0,ytop=1,col="grey",border=NA)
rect(xleft=pred.b$rr.lo,xright=pred.b$rr.up,ybottom=1,ytop=2,col="darkgreen",border=NA)
rect(xleft=pred.c$rr.lo,xright=pred.c$rr.up,ybottom=-1,ytop=0,col="red",border=NA)
segments(x0=pred.a$rr,x1=pred.a$rr,y0=0,y1=1,col="black")
segments(x0=pred.b$rr,x1=pred.b$rr,y0=1,y1=2,col="black")
segments(x0=pred.c$rr,x1=pred.c$rr,y0=0,y1=-1,col="black")
axis(1,at=seq(xlim[1],xlim[2],by=25),cex.axis=cex.axis)
#segments(x0=seq(xlim[1],xlim[2],by=25),x1=seq(xlim[1],xlim[2],by=25),y0=-1,y1=2,lty=3,col="lightblue")
segments(x0=xlim[1],x1=xlim[2],y0=-1:2+.5,y1=-1:2+.5,lty=1,col=c("red","grey","darkgreen"),lwd=.2)
abline(v=0,col="lightblue")}

rr.plot2 <- function(pred.b,pred.c,xlim,cex.axis=.9) {
par(mar=c(2,2,0,0))
plot(x=0,y=0,col=NA,xlim=xlim,ylim=c(-1,1),xaxt="n",yaxt="n",bty="n",las=1)
rect(xleft=pred.b$rr.lo,xright=pred.b$rr.up,ybottom=0,ytop=1,col="darkgreen",border=NA)
rect(xleft=pred.c$rr.lo,xright=pred.c$rr.up,ybottom=-1,ytop=0,col="red",border=NA)
segments(x0=pred.b$rr,x1=pred.b$rr,y0=0,y1=1,col="black")
segments(x0=pred.c$rr,x1=pred.c$rr,y0=0,y1=-1,col="black")
axis(1,at=seq(xlim[1],xlim[2],by=25),cex.axis=cex.axis)
axis(2,at=c(-.5,.5),labels=c("N","I"),las=1,tick=F)
#segments(x0=seq(xlim[1],xlim[2],by=25),x1=seq(xlim[1],xlim[2],by=25),y0=-1,y1=2,lty=3,col="lightblue")
segments(x0=xlim[1],x1=xlim[2],y0=-1:1+.5,y1=-1:1+.5,lty=1,col=c("red","darkgreen"),lwd=.2)
abline(v=0,col="lightblue")}

# 
# rr.plot2 <- function(pred.b,pred.c,xlim,cex.axis=.9) {
#   xmin. <- min(round((min(pred.b$rr.lo,pred.c$rr.lo)-5)/5)*5,-5)
#   xmax. <- max(round((max(pred.b$rr.up,pred.c$rr.up)+5)/5)*5,5)
#   xlim <- c(xmin.,xmax.)
# par(mar=c(2,2,0,0))
# plot(x=0,y=0,col=NA,xlim=xlim,ylim=c(-1,1),xaxt="n",yaxt="n",bty="n",las=1)
# rect(xleft=pred.b$rr.lo,xright=pred.b$rr.up,ybottom=0,ytop=1,col="darkgreen",border=NA)
# rect(xleft=pred.c$rr.lo,xright=pred.c$rr.up,ybottom=-1,ytop=0,col="red",border=NA)
# segments(x0=pred.b$rr,x1=pred.b$rr,y0=0,y1=1,col="black")
# segments(x0=pred.c$rr,x1=pred.c$rr,y0=0,y1=-1,col="black")
# axis(1,at=seq(xlim[1],xlim[2],by=25),cex.axis=cex.axis)
# axis(2,at=c(-.5,.5),labels=c("N","I"),las=1,tick=F)
# #segments(x0=seq(xlim[1],xlim[2],by=25),x1=seq(xlim[1],xlim[2],by=25),y0=-1,y1=2,lty=3,col="lightblue")
# segments(x0=xlim[1],x1=xlim[2],y0=-1:1+.5,y1=-1:1+.5,lty=1,col=c("red","darkgreen"),lwd=.2)
# abline(v=0,col="lightblue")}


rr.plot3 <- function(pred.b,pred.c,cex.axis=.9,xlim=NULL) {
  if(length(xlim)==0){
  xmin. <- min(round((min(pred.b$rr.lo,pred.c$rr.lo)-5)/5)*5,-5)
  xmax. <- max(round((max(pred.b$rr.up,pred.c$rr.up)+5)/5)*5,5)
  xlim <- c(xmin.,xmax.)}
par(mar=c(2,2,0,0))
plot(x=0,y=0,col=NA,xlim=xlim,ylim=c(-1,1),xaxt="n",yaxt="n",bty="n",las=1)
rect(xleft=pred.b$rr.lo,xright=pred.b$rr.up,ybottom=0,ytop=1,col="darkgreen",border=NA)
rect(xleft=pred.c$rr.lo,xright=pred.c$rr.up,ybottom=-1,ytop=0,col="red",border=NA)
segments(x0=pred.b$rr,x1=pred.b$rr,y0=0,y1=1,col="black")
segments(x0=pred.c$rr,x1=pred.c$rr,y0=0,y1=-1,col="black")
axis(1,at=seq(xlim[1],xlim[2],by=5),cex.axis=cex.axis)
axis(2,at=c(-.5,.5),labels=c("N","I"),las=1,tick=F)
#segments(x0=seq(xlim[1],xlim[2],by=25),x1=seq(xlim[1],xlim[2],by=25),y0=-1,y1=2,lty=3,col="lightblue")
segments(x0=xlim[1],x1=xlim[2],y0=-1:1+.5,y1=-1:1+.5,lty=1,col=c("red","darkgreen"),lwd=.2)
abline(v=0,col="lightblue")}

##############################################
## Better GLM predict function
##############################################

#mod <- mod_n2c
#var <- "LANGUAGE"
#vals <- c(50,95)

predman <- function(mod,var,vals){
vars <- unlist(strsplit(as.character(mod$formula)[3],split=" + ",fixed=T))
n <- length(vals)
X <- as.data.frame(matrix(apply(mod$model[,vars],2,function(x){mean(x,na.rm=T)}),nrow=n,ncol=length(vars),byrow=T))
names(X) <- vars
X[,var] <- vals
preds <- predict.glm(mod,newdata=X,type="link",se.fit=T)
pred.y <- preds$fit
pred.lo <- preds$fit-1.96*preds$se.fit
pred.up <- preds$fit+1.96*preds$se.fit
pred.y <- 1/(1+exp(-pred.y))
pred.lo <- 1/(1+exp(-pred.lo))
pred.up <- 1/(1+exp(-pred.up))

y1 <- rnorm(n=10000,mean=preds$fit[length(preds$fit)],sd=preds$se.fit[length(preds$se.fit)])
y0 <- rnorm(n=10000,mean=preds$fit[1],sd=preds$se.fit[1])
pred.y1 <- 1/(1+exp(-y1))
pred.y0 <- 1/(1+exp(-y0))
rrs <- 100*(pred.y1/pred.y0-1)
rr <- mean(rrs)
rr.lo <- quantile(rrs,.025)
rr.up <- quantile(rrs,.975)

#rr <- 100*(pred.y[length(pred.y)]/pred.y[1]-1)
#rr.lo <- 100*(pred.lo[length(pred.lo)]/pred.lo[1]-1)
#rr.up <- 100*(pred.up[length(pred.up)]/pred.up[1]-1)

return(list(y=pred.y,lo=pred.lo,up=pred.up, rr=rr,rr.lo=rr.lo,rr.up=rr.up))}




##############################################
## EVEN Better GLM predict function
##############################################

#mod <- mod_n2c
#var <- "LANGUAGE"
#vals <- c(50,95)

predman2 <- function(mod,var,vals){
  nsim<-1000
  vars <- unlist(strsplit(as.character(mod$formula)[3],split=" + ",fixed=T))
  n <- length(vals)
  X <- as.data.frame(matrix(apply(mod$model[,vars],2,function(x){mean(x,na.rm=T)}),nrow=n,ncol=length(vars),byrow=T))
  X <- cbind(rep(1,n),X)
  names(X) <- names(coefficients(mod))
  X[,var] <- vals
  betas <- rmvnorm(n=nsim,mean=coefficients(mod),sigma=vcov(mod))
  mu <- betas%*%t(X)
  pred.y <- mod$family$linkinv(mu)
  y <- apply(pred.y,2,mean)
  y.lo <- apply(pred.y,2,function(x){quantile(x,.025)})
  y.hi <- apply(pred.y,2,function(x){quantile(x,.975)})
  rrs <- 100*(pred.y[,2]/pred.y[,1]-1)
  rr <- mean(rrs)
  rr.lo <- quantile(rrs,.025)
  rr.up <- quantile(rrs,.975)
  return(list(y=y,lo=y.lo,up=y.hi, rr=rr,rr.lo=rr.lo,rr.up=rr.up))}



#################################################
## Nils' functions
#################################################




streg.sort <- function(data, unitvar, timevar) {
  data[order(data[,unitvar], data[,timevar]),]
}

streg.tlag <- function(data, timevar, lagvar, order=1) {
  result <- numeric(nrow(data))
  timevalues <- data[,timevar]
  varvalues <- data[,lagvar]
  times <- unique(timevalues)
  
  # set initial lag values to NA
  for (time in 1:order) {
    result[timevalues==times[time]] <- NA
  }
  
  # compute lag
  for (time in (order+1):length(times)) {
    result[timevalues==times[time]] <- varvalues[timevalues==times[time-order]]
  }
  result
}

streg.slag <- function(data, timevar, lagvar, Wmat) {
  result <- numeric(nrow(data))
  timevalues <- data[,timevar]
  varvalues <- data[,lagvar]
  times <- unique(timevalues)
  for (time in 1:length(times)) {
    result[timevalues==times[time]] <- Wmat %*% varvalues[timevalues==times[time]]
  }
  result
}






#################################################
## Manual coarsening
#################################################

coarsen <- function(x){

x$DIST_PIPES <- 25*round(x$DIST_PIPES/25)
x$DIST_MIL <- 25*round(x$DIST_MIL/25)
x$CHKINT_NEAR <- 50*round(x$CHKINT_NEAR/50)
x$REFUGEE_MIN <- 50*round(x$REFUGEE_MIN/50)
x$POP <- 5000*round(x$POP/5000)
x$LAT <- round(x$LAT)
x$LONG <- round(x$LONG) 
x$SLOPE <- 10*round(x$SLOPE/10)
x$ELEVATION <- 1000*round(x$ELEVATION/1000)
x$LANGUAGE <- 10*round(x$LANGUAGE/10)

return(x)}











#################################################
## Full CEM results
#################################################

cem.est <- function(dat,trt,t,l,d,k,ms){

trts <- c("FC","SI")# trt
types <- c("I","N")# t
lags <- 1:12# l
defs <- 1:4# d
k2ks <- 0:1# k
methods <- c('euclidean', 'maximum', 'manhattan', 'canberra', 'binary', 'minkowski')
# ms

#dat <- data.frame(matrix(NA,0,25))

predata <- prematch.data(treatment=trts[trt],lag=lags[l],def=defs[d],type=types[t],bin=1,data=data)
rightside <- c("Pre.b","PreG.b","PreW","PreGW","ELEVATION","SLOPE","FOREST","POP","LANGUAGE","REFUGEE_MIN","DIST_PIPES","DIST_MIL","DEPORTED","YEAR","MONTH","LONG","LAT","CHKINT_NEAR")
X <- predata[,c("Post","Treat",rightside)]
X <- na.omit(X)

pars <- data.frame(treatment=trts[trt],type=types[t],def=defs[d],window=lags[l],bin=1,k2k=k2ks[k],method=methods[ms])
mat <- cem(treatment="Treat",data=X,drop="Post",k2k=k2ks[k],method=methods[ms],eval.imbalance=T)
imb <- full.imbalance(X=X,mat=mat)
imb.sum <- data.frame(n.pre=imb$Pre$Summary$N,n.post=imb$Post$Summary$N,sdm.pre=imb$Pre$Summary$meanSDM,sd.post=imb$Post$Summary$meanSDM,l1.pre=imb$Pre$Summary$L1,l1.post=imb$Post$Summary$L1,lcs.pre=imb$Pre$Summary$LCS,lcs.post=imb$Post$Summary$LCS,ks.pre=imb$Pre$Summary$minKSp,ks.post=imb$Post$Summary$minKSp)
attform1 <- as.formula(paste("Post ~ Treat"))
attform2 <- as.formula(paste("Post ~ Treat+",paste(rightside,collapse="+")))
m.X <- X[mat$matched,]
est <- att(mat,attform1,data=X,model="linear-RE")
x0 <- c(1,0)
x1 <- c(1,1)
betas <- est[[1]]["Estimate",]
betas.lo <- est[[1]]["Estimate",]-1.96*est[[1]]["Std. Error",]
betas.hi <- est[[1]]["Estimate",]+1.96*est[[1]]["Std. Error",]
atts1 <- data.frame(att1=est[[1]]["Estimate","Treat"],att1.lo=est[[1]]["Estimate","Treat"]-1.96*est[[1]]["Std. Error","Treat"],att1.up=est[[1]]["Estimate","Treat"]+1.96*est[[1]]["Std. Error","Treat"],rr1=100*(sum(betas*x1)/sum(betas*x0)-1))
est <- att(mat,attform2,data=X,model="linear-RE")
x0 <- c(1,0,apply(m.X[,rightside],2,median))
x1 <- c(1,1,apply(m.X[,rightside],2,median))
betas <- est[[1]]["Estimate",]
betas.lo <- est[[1]]["Estimate",]-1.96*est[[1]]["Std. Error",]
betas.hi <- est[[1]]["Estimate",]+1.96*est[[1]]["Std. Error",]
atts2 <- data.frame(att2=est[[1]]["Estimate","Treat"],att2.lo=est[[1]]["Estimate","Treat"]-1.96*est[[1]]["Std. Error","Treat"],att2.up=est[[1]]["Estimate","Treat"]+1.96*est[[1]]["Std. Error","Treat"],rr2=100*(sum(betas*x1)/sum(betas*x0)-1))
dat <- rbind(dat,cbind(pars,imb.sum,atts1,atts2))
names(dat) <- names(cbind(pars,imb.sum,atts1,atts2))
return(dat)
}


############################
## Imbalance stats
############################

full.imbalance <- function(X,mat){

# Pre
imb <- imbalance(group=X$Treat,data=X,drop=c("Post","Treat"))
imb.mat <- cbind(apply(X[X$Treat==1,],2,mean),apply(X[X$Treat==0,],2,mean))
imb.mat <- as.data.frame(imb.mat)
names(imb.mat) <- c("T","C")
imb.mat$SDM <- (imb.mat$T-imb.mat$C)/sd(imb.mat$T)
imb.mat$ks.pval <- NA
for(i in 1:ncol(X)){imb.mat$ks.pval[i] <- ks.test(X[X$Treat==0,i],X[X$Treat==1,i])$p.value}
imb.mat$L1 <- c(NA,NA,imb$tab$L1)
imb.mat <- imb.mat[-c(1:2),]
pre.tab <- list(Tab=imb.mat,
Summary=list(N=nrow(X),T=sum(X$Treat),C=nrow(X)-sum(X$Treat),meanSDM=mean(abs(imb.mat$SDM)),minKSp=min(imb.mat$ks.pval)
,meanL1=mean(imb.mat$L1),L1=imb$L1$L1,LCS=imb$L1$LCS))

# Post
m.X <- X[mat$matched,]
imb.mat <- cbind(apply(m.X[m.X$Treat==1,],2,mean),apply(m.X[m.X$Treat==0,],2,mean))
imb.mat <- as.data.frame(imb.mat)
names(imb.mat) <- c("T","C")
imb.mat$SDM <- (imb.mat$T-imb.mat$C)/sd(imb.mat$T)
imb.mat$ks.pval <- NA
for(i in 1:ncol(m.X)){imb.mat$ks.pval[i] <- ks.test(m.X[m.X$Treat==0,i],m.X[m.X$Treat==1,i])$p.value}
imb.mat$L1 <- c(NA,NA,mat$imbalance$tab$L1)
imb.mat <- imb.mat[-c(1:2),]
post.tab <- list(Tab=imb.mat,
Summary=list(N=nrow(m.X),T=sum(m.X$Treat),C=nrow(m.X)-sum(m.X$Treat),meanSDM=mean(abs(imb.mat$SDM)),minKSp=min(imb.mat$ks.pval)
,meanL1=mean(imb.mat$L1),L1=mat$imbalance$L1$L1,LCS=mat$imbalance$L1$LCS))
return(list(Pre=pre.tab,Post=post.tab))
}


#############################
## Prepare data for matching
#############################
treatment <- "FC"
lag <- 12
def <- 2
type <- "I"
bin <- 1

prematch.data <- function(treatment,lag,def,type,bin,data){

if(treatment=="FC"){
subdata <- data[data$GOV_ALL_F.b+data$GOV_ALL_L.b>0,]
subdata$Treat <- subdata$GOV_ALL_L.b*(1-subdata$GOV_ALL_F.b)}
if(treatment=="SI"){
subdata <- data[data$GOV_SEL.b+data$GOV_IND.b>0,]
subdata$Treat <- subdata$GOV_SEL.b*(1-subdata$GOV_IND.b)}
if(treatment=="ISL"){
  subdata <- data[data[,paste("INS_ALL_I",def,".b",sep="")]+data[,paste("INS_ALL_N",def,".b",sep="")]>0,]
  subdata$Treat <- subdata[,paste("INS_ALL_I",def,".b",sep="")]*(1-subdata[,paste("INS_ALL_N",def,".b",sep="")])}

if(type!="A"){
subdata$Post <- subdata[,paste("INS_ALL_",type,def,"_p",lag,sep="")]
subdata$PostG <- subdata[,paste("GOV_ALL_p",lag,sep="")]
subdata$Pre <- subdata[,paste("INS_ALL_",type,def,"_t",lag,sep="")]
subdata$PreG <- subdata[,paste("GOV_ALL_t",lag,sep="")]

# Temp:
#subdata$Pre <- subdata[,paste("INS_ALL",def,"_t",lag,sep="")]
subdata$PreG <- subdata[,paste("GOV_ALL_t",lag,sep="")]


if(bin==1){
subdata$Post.b <- as.integer(subdata[,paste("INS_ALL_",type,def,"_p",lag,sep="")]>0)
subdata$PostG.b <- as.integer(subdata[,paste("GOV_ALL_p",lag,sep="")]>0)
subdata$Pre.b <- as.integer(subdata[,paste("INS_ALL_",type,def,"_t",lag,sep="")]>0)
subdata$PreG.b <- as.integer(subdata[,paste("GOV_ALL_t",lag,sep="")]>0)

subdata$Pre.b <- as.integer(subdata[,paste("INS_ALL",def,"_t",lag,sep="")]>0)
subdata$PreG.b <- as.integer(subdata[,paste("GOV_ALL_t",lag,sep="")]>0)
}}

if(type=="A"){
subdata$Post <- subdata[,paste("INS_ALL_",type,def,"_p",lag,sep="")]
subdata$PostG <- subdata[,paste("GOV_ALL_p",lag,sep="")]
subdata$Pre <- subdata[,paste("INS_ALL_",type,def,"_t",lag,sep="")]
subdata$PreG <- subdata[,paste("GOV_ALL_t",lag,sep="")]
if(bin==1){
subdata$Post.b <- as.integer(subdata[,paste("INS_ALL_",type,def,"_p",lag,sep="")]>0)
subdata$PostG.b <- as.integer(subdata[,paste("GOV_ALL_p",lag,sep="")]>0)
subdata$Pre.b <- as.integer(subdata[,paste("INS_ALL_",type,def,"_t",lag,sep="")]>0)
subdata$PreG.b <- as.integer(subdata[,paste("GOV_ALL_t",lag,sep="")]>0)}}


subdata$PreW <- subdata[,paste("W_INS_ALL_t",lag,sep="")]
subdata$PreGW <- subdata[,paste("W_GOV_ALL_t",lag,sep="")]

if(treatment=="ISL"){
subdata$PostG_SEL.b <- 1*(subdata[,paste("GOV_SEL_p",lag,sep="")]>0)*(subdata[,paste("GOV_IND_p",lag,sep="")]==0)
subdata$PostG_SEL_IND <- subdata[,paste("GOV_SEL_p",lag,sep="")] - subdata[,paste("GOV_IND_p",lag,sep="")]
subdata$PostG_SEL <- subdata[,paste("GOV_SEL_p",lag,sep="")]
subdata$PostG_IND <- subdata[,paste("GOV_IND_p",lag,sep="")]
subdata$PreG_SEL <- subdata[,paste("GOV_SEL_t",lag,sep="")]
subdata$PreG_IND <- subdata[,paste("GOV_IND_t",lag,sep="")]

}


#if(bin==1){
#subdata$PreW.b <- as.integer(subdata[,paste("W_INS_ALL.b_t",lag,sep="")]<quantile(subdata[,paste("W_INS_ALL.b_t",lag,sep="")],.1,na.rm=T))
#subdata$PreGW.b <- as.integer(subdata[,paste("W_GOV_ALL.b_t",lag,sep="")]<quantile(subdata[,paste("W_GOV_ALL.b_t",lag,sep="")],.1,na.rm=T))
#}
return(subdata)
}



########################################################
## Function: Coefficients for binomial models

coef.bin <- function(mod,rnd){
mat <- as.data.frame(matrix(NA,nrow(mod$coef),3))
names(mat) <- c("Variable","Estimate","S.E.")
coefs <- c()
ses <- c()
ps <- c()
for(j in 1:nrow(mod$coef)){
coefs[j] <- if(abs(mod$coef[j,1])>5*(1/10)^(rnd+1)){round(mod$coef[j,1],rnd)}else{round(mod$coef[j,1],rnd+1)}
ses[j] <- if(abs(mod$coef[j,2])>5*(1/10)^(rnd+1)){round(mod$coef[j,2],rnd)}else{round(mod$coef[j,2],rnd+1)}
ifelse(mod$coef[j,4]<.001,{ps[j] <- "***"},
ifelse(mod$coef[j,4]<.01&mod$coef[j,4]>=.001,{ps[j] <- "**"},
ifelse(mod$coef[j,4]<.05&mod$coef[j,4]>=.01,{ps[j] <- "*"},"")))
}
mat[,1:2] <- cbind(rownames(mod$coef),coefs)
mat[,3] <- paste("(",ses,")",ps,sep="")

return(mat)}

#######################################################
## Function: Coefficients for count models

coef.count <- function(mod,rnd){
mat1 <- as.data.frame(matrix(NA,nrow(mod$coef$count),3))
names(mat1) <- c("Variable","Estimate","S.E.")
coefs <- c()
ses <- c()
ps <- c()


for(j in 1:nrow(mod$coef$count)){
if(!is.na(mod$coef$count[j,1])){
coefs[j] <- if(abs(mod$coef$count[j,1])>5*(1/10)^(rnd+1)){round(mod$coef$count[j,1],rnd)}else{round(mod$coef$count[j,1],rnd+1)}}
if(!is.na(mod$coef$count[j,2])){
ses[j] <- if(abs(mod$coef$count[j,2])>5*(1/10)^(rnd+1)){round(mod$coef$count[j,2],rnd)}else{round(mod$coef$count[j,2],rnd+1)}}
if(!is.na(mod$coef$count[j,4])){
ifelse(mod$coef$count[j,4]<.001,{ps[j] <- "***"},
ifelse(mod$coef$count[j,4]<.01&mod$coef$count[j,4]>=.001,{ps[j] <- "**"},
ifelse(mod$coef$count[j,4]<.05&mod$coef$count[j,4]>=.01,{ps[j] <- "*"},"")))
}}
mat1[,1:2] <- cbind(rownames(mod$coef$count),coefs)
mat1[,3] <- paste("(",ses,")",ps,sep="")

mat2 <- as.data.frame(matrix(NA,nrow(mod$coef$zero),3))
names(mat2) <- c("Variable","Estimate","S.E.")
coefs <- c()
ses <- c()
ps <- c()

for(j in 1:nrow(mod$coef$zero)){
if(!is.na(mod$coef$zero[j,1])){
coefs[j] <- if(abs(mod$coef$zero[j,1])>5*(1/10)^(rnd+1)){round(mod$coef$zero[j,1],rnd)}else{round(mod$coef$zero[j,1],rnd+1)}}
if(!is.na(mod$coef$zero[j,2])){
ses[j] <- if(abs(mod$coef$zero[j,2])>5*(1/10)^(rnd+1)){round(mod$coef$zero[j,2],rnd)}else{round(mod$coef$zero[j,2],rnd+1)}}
if(!is.na(mod$coef$zero[j,4])){
ifelse(mod$coef$zero[j,4]<.001,{ps[j] <- "***"},
ifelse(mod$coef$zero[j,4]<.01&mod$coef$zero[j,4]>=.001,{ps[j] <- "**"},
ifelse(mod$coef$zero[j,4]<.05&mod$coef$zero[j,4]>=.01,{ps[j] <- "*"},"")))
}}

mat2[,1:2] <- cbind(rownames(mod$coef$zero),coefs)
mat2[,3] <- paste("(",ses,")",ps,sep="")
return(list(count=mat1,zero=mat2))}



##############################################
## Force into 1 column: binomial

coef1.bin <- function(mod,rnd){
mat <- coef.bin(mod,rnd)

mat1 <- as.data.frame(matrix(NA,nrow(mat)*2,2))
names(mat1) <- c("Variable","Estimate (SE)")

odds <- grepl(".5",1:nrow(mat1)/2)
evens <- !grepl(".5",1:nrow(mat1)/2)
mat1[odds,1] <- mat[,1]
mat1[evens,1] <- ""
mat1[odds,2] <- mat[,2]
mat1[evens,2] <- mat[,3]
return(mat1)}


##############################################
## Force into 1 column: count

coef1.count <- function(mod,rnd){
mat <- coef.count(mod,rnd)$count
mat1 <- as.data.frame(matrix(NA,nrow(mat)*2,2))
names(mat1) <- c("Variable","Estimate (SE)")

odds <- grepl(".5",1:nrow(mat1)/2)
evens <- !grepl(".5",1:nrow(mat1)/2)
mat1[odds,1] <- mat[,1]
mat1[evens,1] <- ""
mat1[odds,2] <- mat[,2]
mat1[evens,2] <- mat[,3]


mat <- coef.count(mod,rnd)$zero
mat2 <- as.data.frame(matrix(NA,nrow(mat)*2,2))
names(mat2) <- c("Variable","Estimate (SE)")

odds <- grepl(".5",1:nrow(mat2)/2)
evens <- !grepl(".5",1:nrow(mat2)/2)
mat2[odds,1] <- mat[,1]
mat2[evens,1] <- ""
mat2[odds,2] <- mat[,2]
mat2[evens,2] <- mat[,3]

return(list(count=mat1,zero=mat2))}


###########################################
## AUC

AUC <- function(y, fitted){
    n1 <- sum(y)
    n <- length(y)
    out <- (mean(rank(fitted)[y == 1]) - (n1 + 1)/2)/(n - n1)
    return(out)
}

###########################################
## D-in-D
###########################################


dind <- function(x,post,pre,treatment,rnd,ci){
	

y1c <- mean(x[,post][x[,treatment]==0],na.rm=T)
y0c <- mean(x[,pre][x[,treatment]==0],na.rm=T)
y1t <- mean(x[,post][x[,treatment]==1],na.rm=T)
y0t <- mean(x[,pre][x[,treatment]==1],na.rm=T)
y1c.lo <- quantile(x[,post][x[,treatment]==0],.025,na.rm=T)
y0c.lo <- quantile(x[,pre][x[,treatment]==0],.025,na.rm=T)
y1t.lo <- quantile(x[,post][x[,treatment]==1],.025,na.rm=T)
y0t.lo <- quantile(x[,pre][x[,treatment]==1],.025,na.rm=T)
y1c.hi <- quantile(x[,post][x[,treatment]==0],.975,na.rm=T)
y0c.hi <- quantile(x[,pre][x[,treatment]==0],.975,na.rm=T)
y1t.hi <- quantile(x[,post][x[,treatment]==1],.975,na.rm=T)
y0t.hi <- quantile(x[,pre][x[,treatment]==1],.975,na.rm=T)
yc <- y1c-y0c
yt <- y1t-y0t
yc.lo <- y1c.lo-y0c.lo
yt.lo <- y1t.lo-y0t.lo
yc.hi <- y1c.hi-y0c.hi
yt.hi <- y1t.hi-y0t.hi
dind <- (yt)-(yc)
dind.lo <- (yt.lo)-(yc.lo)
dind.hi <- (yt.hi)-(yc.hi)
rr.c <- y1c/y0c
rr.t <- y1t/y0t
rr.c.lo <- y1c.lo/y0c.lo
rr.t.lo <- y1t.lo/y0t.lo
rr.c.hi <- y1c.hi/y0c.hi
rr.t.hi <- y1t.hi/y0t.hi
rr100.c <- 100*(rr.c-1)
rr100.t <- 100*(rr.t-1)
rr100.c.lo <- 100*(rr.c.lo-1)
rr100.t.lo <- 100*(rr.t.lo-1)
rr100.c.hi <- 100*(rr.c.hi-1)
rr100.t.hi <- 100*(rr.t.hi-1)
rr100.d <- rr100.t-rr100.c
rr100.d.lo <- rr100.t.lo-rr100.c.lo
rr100.d.hi <- rr100.t.lo-rr100.c.hi

y1c <- round(y1c,rnd)
y0c <- round(y0c,rnd)
y1t <- round(y1t,rnd)
y0t <- round(y0t,rnd)
y1c.lo <- round(y1c.lo,rnd)
y0c.lo <- round(y0c.lo,rnd)
y1t.lo <- round(y1t.lo,rnd)
y0t.lo <- round(y0t.lo,rnd)
y1c.hi <- round(y1c.hi,rnd)
y0c.hi <- round(y0c.hi,rnd)
y1t.hi <- round(y1t.hi,rnd)
y0t.hi <- round(y0t.hi,rnd)
yc <- round(yc,rnd)
yt <- round(yt,rnd)
yc.lo <- round(yc.lo,rnd)
yt.lo <- round(yt.lo,rnd)
yc.hi <- round(yc.hi,rnd)
yt.hi <- round(yt.hi,rnd)
dind <- round(dind,rnd)
dind.lo <- round(dind.lo,rnd)
dind.hi <- round(dind.hi,rnd)
rr.c <- round(rr.c,rnd)
rr.t <- round(rr.t,rnd)
rr.c.lo <- round(rr.c.lo,rnd)
rr.t.lo <- round(rr.t.lo,rnd)
rr.c.hi <- round(rr.c.hi,rnd)
rr.t.hi <- round(rr.t.hi,rnd)
rr100.c <- round(rr100.c,rnd)
rr100.t <- round(rr100.t,rnd)
rr100.c.lo <- round(rr100.c.lo,rnd)
rr100.t.lo <- round(rr100.t.lo,rnd)
rr100.c.hi <- round(rr100.c.hi,rnd)
rr100.t.hi <- round(rr100.t.hi,rnd)
rr100.d <- round(rr100.d,rnd)
rr100.d.lo <- round(rr100.d.lo,rnd)
rr100.d.hi <- round(rr100.d.hi,rnd)

dind.mat <- as.data.frame(matrix(NA,5,4))
names(dind.mat) <- c("Quantity","Control","Treatment","Diff-in-Diff")
dind.mat[,1] <- c("Y(0)","Y(1)","Y(1)-Y(0)","Y(1)/Y(0)","Y(1)/Y(0) %")
dind.mat[,2] <- c(y0c,y1c,yc,rr.c,rr100.c)
dind.mat[,3] <- c(y0t,y1t,yt,rr.t,rr100.t)
dind.mat[1,4] <- y0t-y0c
dind.mat[2,4] <- y1t-y1c
dind.mat[3,4] <- dind
dind.mat[4,4] <- rr.t-rr.c
dind.mat[5,4] <- rr100.d

dind.mat.lo <- as.data.frame(matrix(NA,5,4))
names(dind.mat.lo) <- c("Quantity","Control","Treatment","Diff-in-Diff")
dind.mat.lo[,1] <- c("Y(0)","Y(1)","Y(1)-Y(0)","Y(1)/Y(0)","Y(1)/Y(0) %")
dind.mat.lo[,2] <- c(y0c.lo,y1c.lo,yc.lo,rr.c.lo,rr100.c.lo)
dind.mat.lo[,3] <- c(y0t.lo,y1t.lo,yt.lo,rr.t.lo,rr100.t.lo)
dind.mat.lo[1,4] <- y0t.lo-y0c.lo
dind.mat.lo[2,4] <- y1t.lo-y1c.lo
dind.mat.lo[3,4] <- dind.lo
dind.mat.lo[4,4] <- rr.t.lo-rr.c.lo
dind.mat.lo[5,4] <- rr100.d.lo

dind.mat.hi <- as.data.frame(matrix(NA,5,4))
names(dind.mat.hi) <- c("Quantity","Control","Treatment","Diff-in-Diff")
dind.mat.hi[,1] <- c("Y(0)","Y(1)","Y(1)-Y(0)","Y(1)/Y(0)","Y(1)/Y(0) %")
dind.mat.hi[,2] <- c(y0c.hi,y1c.hi,yc.hi,rr.c.hi,rr100.c.hi)
dind.mat.hi[,3] <- c(y0t.hi,y1t.hi,yt.hi,rr.t.hi,rr100.t.hi)
dind.mat.hi[1,4] <- y0t.hi-y0c.hi
dind.mat.hi[2,4] <- y1t.hi-y1c.hi
dind.mat.hi[3,4] <- dind.hi
dind.mat.hi[4,4] <- rr.t.hi-rr.c.hi
dind.mat.hi[5,4] <- rr100.d.hi

if(ci==T){return(list(Diff=dind.mat,Lower=dind.mat.lo,Upper=dind.mat.hi))}
if(ci==F){dind.mat}
}






###############################
## Correlation matrix plot
###############################


cormat <- function(data,vars,labels){

Xs.main <- data[,vars]
Xs.main <- na.omit(Xs.main)

mat.cor <- round(cor(Xs.main),2)
rownames(mat.cor) <- labels
colnames(mat.cor) <- rownames(mat.cor) 

# Color palette
pal <- c("darkblue","blue","lightblue","lightblue1","grey95","beige","lightgoldenrod1","orange","red","darkred")
cols.mat <- matrix(NA,nrow=nrow(mat.cor),ncol=ncol(mat.cor))
cols.mat <- ifelse(mat.cor>=-1&mat.cor<(-.8),pal[1],
ifelse(mat.cor>=-(.8)&mat.cor<(-.6),pal[2],
ifelse(mat.cor>=-(.6)&mat.cor<(-.4),pal[3],
ifelse(mat.cor>=-(.4)&mat.cor<(-.2),pal[4],
ifelse(mat.cor>=-(.2)&mat.cor<(0),pal[5],
ifelse(mat.cor>=-(0)&mat.cor<(.2),pal[6],
ifelse(mat.cor>=-(.2)&mat.cor<(.4),pal[7],
ifelse(mat.cor>=-(.4)&mat.cor<(.6),pal[8],
ifelse(mat.cor>=-(.6)&mat.cor<(.8),pal[9],
pal[10]
)))))))))
cols.mat <- cols.mat[ nrow(cols.mat):1,]

# pch matrix
pchs.mat <- matrix(NA,nrow=nrow(mat.cor),ncol=ncol(mat.cor))
pchs.mat <- ifelse(mat.cor<0,4,NA)
pchs.mat <- pchs.mat[ nrow(pchs.mat):1,]


par(mar=c(0,0,0,0))
plot(c(-8,nrow(mat.cor)+1), c(0,nrow(mat.cor)+3), type = "n", xlab="", ylab="", asp = 1,axes=F)

ups <- upper.tri(cols.mat,diag=T)[ncol(cols.mat):1,]
cols.mat[which(ups)] <- NA
ups <- upper.tri(pchs.mat,diag=T)[ncol(pchs.mat):1,]
pchs.mat[which(ups)] <- NA

for(i in 1:nrow(cols.mat)){
for(j in 1:ncol(cols.mat)){
points(x=i,y=j,cex=2.4,col=cols.mat[j,i],pch=15)
points(x=i,y=j,cex=2.4,col="grey",pch=pchs.mat[j,i])
}
}
text(x=0.5,y=ncol(cols.mat):1,pos=2,labels = c(NA,labels[-1]),cex=.7)
text(x=1:(ncol(cols.mat))-.35,y=ncol(cols.mat):1,pos=4,labels = c(labels[-length(labels)],NA),cex=.75, srt = 90)

legend.v2(x="bottom",cex=.8,fill=pal,legend=c("[-1, -.8)","[-.8, -.6)","[-.6, -.4)","[-.4, -.2)","[-.2, 0)","[0, .2)","[.2, .4)","[.4, .6)","[.6, .8)","[.8, 1]"),title="Correlation",ncol=5,pch=c(4,4,4,4,4,NA,NA,NA,NA,NA),bty="n",col=c(rep("darkgrey",5),rep(NA,5)))



}


###############################
## Correlation matrix legend
###############################

legend.v2 <- function (x, y = NULL, legend, fill = NULL, col = par("col"), border = "black", lty, lwd, pch, angle = 45, density = NULL, bty = "o", bg = par("bg"), box.lwd = par("lwd"), box.lty = par("lty"), box.col = par("fg"), pt.bg = NA, cex = 1, pt.cex = cex, pt.lwd = lwd, xjust = 0, yjust = 1, x.intersp = 1, y.intersp = 1, adj = c(0, 0.5), text.width = NULL, text.col = par("col"), merge = do.lines && has.pch, trace = FALSE, plot = TRUE, ncol = 1, horiz = FALSE, title = NULL, inset = 0, xpd, title.col = text.col, title.adj = 0.5,seg.len = 2){
if (missing(legend) && !missing(y) && (is.character(y) || 
is.expression(y))) {
legend <- y
y <- NULL
}
mfill <- !missing(fill) || !missing(density)
if (!missing(xpd)) {
op <- par("xpd")
on.exit(par(xpd = op))
par(xpd = xpd)
}
title <- as.graphicsAnnot(title)
if (length(title) > 1) 
stop("invalid title")
legend <- as.graphicsAnnot(legend)
n.leg <- if (is.call(legend)) 
1
else length(legend)
if (n.leg == 0) 
stop("'legend' is of length 0")
auto <- if (is.character(x)) 
match.arg(x, c("bottomright", "bottom", "bottomleft", 
"left", "topleft", "top", "topright", "right", "center"))
else NA
if (is.na(auto)) {
xy <- xy.coords(x, y)
x <- xy$x
y <- xy$y
nx <- length(x)
if (nx < 1 || nx > 2) 
stop("invalid coordinate lengths")
}
else nx <- 0
xlog <- par("xlog")
ylog <- par("ylog")
rect2 <- function(left, top, dx, dy, density = NULL, angle, 
...) {
r <- left + dx
if (xlog) {
left <- 10^left
r <- 10^r
}
b <- top - dy
if (ylog) {
top <- 10^top
b <- 10^b
}
rect(left, top, r, b, angle = angle, density = density, 
...)
}
segments2 <- function(x1, y1, dx, dy, ...) {
x2 <- x1 + dx
if (xlog) {
x1 <- 10^x1
x2 <- 10^x2
}
y2 <- y1 + dy
if (ylog) {
y1 <- 10^y1
y2 <- 10^y2
}
segments(x1, y1, x2, y2, ...)
}
points2 <- function(x, y, ...) {
if (xlog) 
x <- 10^x
if (ylog) 
y <- 10^y
points(x, y, ...)
}
text2 <- function(x, y, ...) {
if (xlog) 
x <- 10^x
if (ylog) 
y <- 10^y
text(x, y, ...)
}
if (trace) 
catn <- function(...) do.call("cat", c(lapply(list(...), 
formatC), list("\n")))
cin <- par("cin")
Cex <- cex * par("cex")
if (is.null(text.width)) 
text.width <- max(abs(strwidth(legend, units = "user", 
cex = cex)))
else if (!is.numeric(text.width) || text.width < 0) 
stop("'text.width' must be numeric, >= 0")
xc <- Cex * xinch(cin[1L], warn.log = FALSE)
yc <- Cex * yinch(cin[2L], warn.log = FALSE)
if (xc < 0) 
text.width <- -text.width
xchar <- xc
xextra <- 0
yextra <- yc * (y.intersp - 1)
ymax <- yc * max(1, strheight(legend, units = "user", cex = cex)/yc)
ychar <- yextra + ymax
if (trace) 
catn("  xchar=", xchar, "; (yextra,ychar)=", c(yextra, 
ychar))
if (mfill) {
xbox <- xc * 0.8
ybox <- yc * 0.5
dx.fill <- xbox
}
do.lines <- (!missing(lty) && (is.character(lty) || any(lty > 
0))) || !missing(lwd)
n.legpercol <- if (horiz) {
if (ncol != 1) 
warning("horizontal specification overrides: Number of columns := ", 
n.leg)
ncol <- n.leg
1
}
else ceiling(n.leg/ncol)
has.pch <- !missing(pch) && length(pch) > 0
if (do.lines) {
x.off <- if (merge) 
-0.7
else 0
}
else if (merge) 
warning("'merge = TRUE' has no effect when no line segments are drawn")
if (has.pch) {
if (is.character(pch) && !is.na(pch[1L]) && nchar(pch[1L], 
type = "c") > 1) {
if (length(pch) > 1) 
warning("not using pch[2..] since pch[1L] has multiple chars")
np <- nchar(pch[1L], type = "c")
pch <- substr(rep.int(pch[1L], np), 1L:np, 1L:np)
}
}
if (is.na(auto)) {
if (xlog) 
x <- log10(x)
if (ylog) 
y <- log10(y)
}
if (nx == 2) {
x <- sort(x)
y <- sort(y)
left <- x[1L]
top <- y[2L]
w <- diff(x)
h <- diff(y)
w0 <- w/ncol
x <- mean(x)
y <- mean(y)
if (missing(xjust)) 
xjust <- 0.5
if (missing(yjust)) 
yjust <- 0.5
}
else {
h <- (n.legpercol + (!is.null(title))) * ychar + yc
w0 <- text.width + (x.intersp + 1) * xchar
if (mfill) 
w0 <- w0 + dx.fill
if (do.lines) 
w0 <- w0 + (seg.len + +x.off) * xchar
w <- ncol * w0 + 0.5 * xchar
if (!is.null(title) && (abs(tw <- strwidth(title, units = "user", 
cex = cex) + 0.5 * xchar)) > abs(w)) {
xextra <- (tw - w)/2
w <- tw
}
if (is.na(auto)) {
left <- x - xjust * w
top <- y + (1 - yjust) * h
}
else {
usr <- par("usr")
inset <- rep(inset, length.out = 2)
insetx <- inset[1L] * (usr[2L] - usr[1L])
left <- switch(auto, bottomright = , topright = , 
right = usr[2L] - w - insetx, bottomleft = , 
left = , topleft = usr[1L] + insetx, bottom = , 
top = , center = (usr[1L] + usr[2L] - w)/2)
insety <- inset[2L] * (usr[4L] - usr[3L])
top <- switch(auto, bottomright = , bottom = , bottomleft = usr[3L] + 
h + insety, topleft = , top = , topright = usr[4L] - 
insety, left = , right = , center = (usr[3L] + 
usr[4L] + h)/2)
}
}
if (plot && bty != "n") {
if (trace) 
catn("  rect2(", left, ",", top, ", w=", w, ", h=", 
h, ", ...)", sep = "")
rect2(left, top, dx = w, dy = h, col = bg, density = NULL, 
lwd = box.lwd, lty = box.lty, border = box.col)
}
xt <- left + xchar + xextra + (w0 * rep.int(0:(ncol - 1), 
rep.int(n.legpercol, ncol)))[1L:n.leg]
yt <- top - 0.5 * yextra - ymax - (rep.int(1L:n.legpercol, 
ncol)[1L:n.leg] - 1 + (!is.null(title))) * ychar
if (mfill) {
if (plot) {
fill <- rep(fill, length.out = n.leg)
rect2(left = xt+.22, top = yt + ybox/2, dx = xbox, dy = ybox, 
col = fill, density = density, angle = angle, 
border = border)
}
xt <- xt + dx.fill
}
if (plot && (has.pch || do.lines)) 
col <- rep(col, length.out = n.leg)
if (missing(lwd)) 
lwd <- par("lwd")
if (do.lines) {
if (missing(lty)) 
lty <- 1
lty <- rep(lty, length.out = n.leg)
lwd <- rep(lwd, length.out = n.leg)
ok.l <- !is.na(lty) & (is.character(lty) | lty > 0)
if (trace) 
catn("  segments2(", xt[ok.l] + x.off * xchar, ",", 
yt[ok.l], ", dx=", seg.len * xchar, ", dy=0, ...)")
if (plot) 
segments2(xt[ok.l] + x.off * xchar, yt[ok.l], dx = seg.len * 
xchar, dy = 0, lty = lty[ok.l], lwd = lwd[ok.l], 
col = col[ok.l])
xt <- xt + (seg.len + x.off) * xchar
}
if (has.pch) {
pch <- rep(pch, length.out = n.leg)
pt.bg <- rep(pt.bg, length.out = n.leg)
pt.cex <- rep(pt.cex, length.out = n.leg)
pt.lwd <- rep(pt.lwd, length.out = n.leg)
ok <- !is.na(pch) & (is.character(pch) | pch >= 0)
x1 <- (if (merge && do.lines) 
xt - (seg.len/2) * xchar
else xt)[ok]
y1 <- yt[ok]
if (trace) 
catn("  points2(", x1, ",", y1, ", pch=", pch[ok], 
", ...)")
if (plot) 
points2(x1, y1, pch = pch[ok], col = col[ok], cex = pt.cex[ok], 
bg = pt.bg[ok], lwd = pt.lwd[ok])
}
xt <- xt + x.intersp * xchar
if (plot) {
if (!is.null(title)) 
text2(left + w * title.adj, top - ymax, labels = title, 
adj = c(title.adj, 0), cex = cex, col = title.col)
text2(xt, yt, labels = legend, adj = adj, cex = cex, 
col = text.col)
}
invisible(list(rect = list(w = w, h = h, left = left, top = top), 
text = list(x = xt, y = yt)))
}

